"""Add InstanceModel.last_processed_at

Revision ID: 7b24b1c8eba7
Revises: ea60480f82bb
Create Date: 2024-08-30 13:42:08.961792

"""

import uuid

import sqlalchemy as sa
import sqlalchemy_utils
from alembic import op

import dstack._internal.server.models
from dstack._internal.utils.common import get_current_datetime

# revision identifiers, used by Alembic.
revision = "7b24b1c8eba7"
down_revision = "ea60480f82bb"
branch_labels = None
depends_on = None


instance_table = sa.Table(
    "instances",
    sa.MetaData(),
    # partial description - only columns affected by this migration
    sa.Column("id", sqlalchemy_utils.UUIDType(binary=False), primary_key=True, default=uuid.uuid4),
    sa.Column("last_processed_at", dstack._internal.server.models.NaiveDateTime()),
)


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("instances", schema=None) as batch_op:
        batch_op.add_column(
            sa.Column(
                "last_processed_at", dstack._internal.server.models.NaiveDateTime(), nullable=True
            )
        )

    op.get_bind().execute(
        (instance_table.update().values(last_processed_at=get_current_datetime()))
    )

    with op.batch_alter_table("instances", schema=None) as batch_op:
        batch_op.alter_column("last_processed_at", nullable=False)

    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.add_column(sa.Column("instance_assigned", sa.Boolean(), nullable=True))

    op.execute(sa.sql.text("UPDATE jobs SET instance_assigned = FALSE"))

    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.alter_column("instance_assigned", nullable=False)
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    with op.batch_alter_table("jobs", schema=None) as batch_op:
        batch_op.drop_column("instance_assigned")

    with op.batch_alter_table("instances", schema=None) as batch_op:
        batch_op.drop_column("last_processed_at")

    # ### end Alembic commands ###
